from flask import Flask, request, jsonify
import random
import string
import requests
import sys
import os
import logging
from dotenv import load_dotenv
from functools import wraps
from datetime import datetime

# 加载环境变量
load_dotenv()

# 配置日志
logging.basicConfig(
    level=os.getenv('LOG_LEVEL', 'INFO'),
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler(os.getenv('LOG_FILE', 'server.log')),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

app = Flask(__name__)

# 配置类
class Config:
    OLLAMA_API_URL = os.getenv('OLLAMA_API_URL', 'http://localhost:11434')
    DEFAULT_PORT = int(os.getenv('DEFAULT_PORT', '4433'))
    TOKEN_LENGTH = int(os.getenv('TOKEN_LENGTH', '32'))
    ALLOWED_ORIGINS = os.getenv('ALLOWED_ORIGINS', '*').split(',')
    RATE_LIMIT_PER_DAY = int(os.getenv('RATE_LIMIT_PER_DAY', '200'))
    RATE_LIMIT_PER_HOUR = int(os.getenv('RATE_LIMIT_PER_HOUR', '50'))
    HEALTH_CHECK_INTERVAL = int(os.getenv('HEALTH_CHECK_INTERVAL', '60'))
    HEALTH_CHECK_TIMEOUT = int(os.getenv('HEALTH_CHECK_TIMEOUT', '5'))

# 生成随机token的函数
def generate_random_token(length=Config.TOKEN_LENGTH):
    charset = string.ascii_letters + string.digits
    return ''.join(random.choice(charset) for _ in range(length))

TOKEN = generate_random_token()
logger.info(f"Generated Token: {TOKEN}")

# 错误处理装饰器
def handle_errors(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        try:
            return f(*args, **kwargs)
        except requests.exceptions.RequestException as e:
            logger.error(f"API request failed: {str(e)}")
            return jsonify({"error": f"Failed to call Ollama API: {str(e)}"}), 500
        except Exception as e:
            logger.error(f"Internal server error: {str(e)}")
            return jsonify({"error": f"Internal server error: {str(e)}"}), 500
    return wrapper

# 认证装饰器
def require_auth(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        auth_header = request.headers.get('Authorization')
        if not auth_header or auth_header != TOKEN:
            logger.warning(f"Unauthorized access attempt from {request.remote_addr}")
            return jsonify({"error": "Unauthorized"}), 401
        return f(*args, **kwargs)
    return wrapper

# 健康检查端点
@app.route('/health', methods=['GET'])
def health_check():
    try:
        # 检查 Ollama API 是否可访问
        response = requests.get(f"{Config.OLLAMA_API_URL}/api/tags", timeout=Config.HEALTH_CHECK_TIMEOUT)
        response.raise_for_status()
        return jsonify({
            "status": "healthy",
            "ollama_api": "available",
            "timestamp": datetime.now().isoformat()
        }), 200
    except Exception as e:
        logger.error(f"Health check failed: {str(e)}")
        return jsonify({
            "status": "unhealthy",
            "error": str(e),
            "timestamp": datetime.now().isoformat()
        }), 503

@app.route('/api/generate', methods=['POST'])
@require_auth
@handle_errors
def generate_text():
    data = request.json
    model = data.get("model")
    prompt = data.get("prompt")
    stream = data.get("stream", False)
    options = data.get("options", {})

    if not model or not prompt:
        logger.warning("Missing required parameters in generate request")
        return jsonify({"error": "Missing required parameters"}), 400

    logger.info(f"Generating text with model: {model}")
    ollama_url = f"{Config.OLLAMA_API_URL}/api/generate"
    payload = {
        "model": model,
        "prompt": prompt,
        "stream": stream,
        "options": options
    }

    response = requests.post(ollama_url, json=payload)
    response.raise_for_status()
    return response.json(), 200

@app.route('/api/chat', methods=['POST'])
@require_auth
@handle_errors
def chat():
    data = request.json
    model = data.get("model")
    messages = data.get("messages")
    stream = data.get("stream", False)
    options = data.get("options", {})

    if not model or not messages:
        logger.warning("Missing required parameters in chat request")
        return jsonify({"error": "Missing required parameters"}), 400

    logger.info(f"Chatting with model: {model}")
    ollama_url = f"{Config.OLLAMA_API_URL}/api/chat"
    payload = {
        "model": model,
        "messages": messages,
        "stream": stream,
        "options": options
    }

    response = requests.post(ollama_url, json=payload)
    response.raise_for_status()
    return response.json(), 200

@app.route('/api/embeddings', methods=['POST'])
@require_auth
@handle_errors
def embeddings():
    data = request.json
    model_name = data.get("model", {}).get("name")
    input_data = data.get("input", [])

    if not model_name or not input_data:
        logger.warning("Missing required parameters in embeddings request")
        return jsonify({"error": "Missing required parameters"}), 400

    logger.info(f"Generating embeddings with model: {model_name}")
    ollama_url = f"{Config.OLLAMA_API_URL}/api/embeddings"
    payload = {
        "model": {
            "name": model_name
        },
        "input": input_data
    }

    response = requests.post(ollama_url, json=payload)
    response.raise_for_status()
    return response.json(), 200

if __name__ == '__main__':
    port = Config.DEFAULT_PORT
    if len(sys.argv) > 1:
        port = int(sys.argv[1])

    logger.info(f"Starting server on port {port}")
    logger.info(f"Ollama API URL: {Config.OLLAMA_API_URL}")
    logger.info(f"Allowed origins: {Config.ALLOWED_ORIGINS}")
    
    app.run(host='0.0.0.0', port=port)